import os
import sys
import numpy as np
from torch.utils.data import Dataset




def load_HAR_dataset(root_path,
                     data_path,
                     data_name):
    data_dir = os.path.join(root_path, data_path, data_name)
    subject_data = np.loadtxt(f'{data_dir}/train/subject_train.txt')
    # Samples
    train_acc_x = np.loadtxt(f'{data_dir}/train/Inertial Signals/body_acc_x_train.txt')
    train_acc_y = np.loadtxt(f'{data_dir}/train/Inertial Signals/body_acc_y_train.txt')
    train_acc_z = np.loadtxt(f'{data_dir}/train/Inertial Signals/body_acc_z_train.txt')
    train_gyro_x = np.loadtxt(f'{data_dir}/train/Inertial Signals/body_gyro_x_train.txt')
    train_gyro_y = np.loadtxt(f'{data_dir}/train/Inertial Signals/body_gyro_y_train.txt')
    train_gyro_z = np.loadtxt(f'{data_dir}/train/Inertial Signals/body_gyro_z_train.txt')
    train_tot_acc_x = np.loadtxt(f'{data_dir}/train/Inertial Signals/total_acc_x_train.txt')
    train_tot_acc_y = np.loadtxt(f'{data_dir}/train/Inertial Signals/total_acc_y_train.txt')
    train_tot_acc_z = np.loadtxt(f'{data_dir}/train/Inertial Signals/total_acc_z_train.txt')

    test_acc_x = np.loadtxt(f'{data_dir}/test/Inertial Signals/body_acc_x_test.txt')
    test_acc_y = np.loadtxt(f'{data_dir}/test/Inertial Signals/body_acc_y_test.txt')
    test_acc_z = np.loadtxt(f'{data_dir}/test/Inertial Signals/body_acc_z_test.txt')
    test_gyro_x = np.loadtxt(f'{data_dir}/test/Inertial Signals/body_gyro_x_test.txt')
    test_gyro_y = np.loadtxt(f'{data_dir}/test/Inertial Signals/body_gyro_y_test.txt')
    test_gyro_z = np.loadtxt(f'{data_dir}/test/Inertial Signals/body_gyro_z_test.txt')
    test_tot_acc_x = np.loadtxt(f'{data_dir}/test/Inertial Signals/total_acc_x_test.txt')
    test_tot_acc_y = np.loadtxt(f'{data_dir}/test/Inertial Signals/total_acc_y_test.txt')
    test_tot_acc_z = np.loadtxt(f'{data_dir}/test/Inertial Signals/total_acc_z_test.txt')

    # Stacking channels together data
    train_data = np.stack((train_acc_x, train_acc_y, train_acc_z,
                           train_gyro_x, train_gyro_y, train_gyro_z,
                           train_tot_acc_x, train_tot_acc_y, train_tot_acc_z), axis=1)
    test_data = np.stack((test_acc_x, test_acc_y, test_acc_z,
                       test_gyro_x, test_gyro_y, test_gyro_z,
                       test_tot_acc_x, test_tot_acc_y, test_tot_acc_z), axis=1)

    # labels
    train_labels = np.loadtxt(f'{data_dir}/train/y_train.txt')
    train_labels -= np.min(train_labels)
    test_labels = np.loadtxt(f'{data_dir}/test/y_test.txt')
    test_labels -= np.min(test_labels)

    return train_data, train_labels, test_data, test_labels


class Dataset_HAR(Dataset):
    def __init__(self,
                 root_path,
                 data_path,
                 data_name,
                 flag='TRAIN',
                 config=None
                ):

        # init
        self.root_path = root_path
        self.data_path = data_path
        self.data_name = data_name
        self.flag = flag

        self.config = config

        self.__read_data()

    def __read_data(self):
        if self.flag == 'TRAIN':
            #print('Loading TRAIN data')
            self.x, self.y, *self._ = load_HAR_dataset(
                root_path=self.root_path,
                data_path=self.data_path,
                data_name=self.data_name)
        elif self.flag == 'TEST':
            #print('Loading TEST data')
            *self._, self.x, self.y = load_HAR_dataset(
                root_path=self.root_path,
                data_path=self.data_path,
                data_name=self.data_name
            )
        else:
            raise ValueError('Unknown flag, TRAIN or TEST is required')

    def __getitem__(self, idx):
        seq_x = self.x[idx].T
        seq_y = self.y[idx]
        return seq_x, seq_y

    def __len__(self):
        return self.x.shape[0]


if __name__ == "__main__":

    root_path = os.path.abspath('../data')
    data_path = 'HAR'
    data_name = 'UCI HAR Dataset'

    train_data, train_labels, test_data, test_labels = load_HAR_dataset(
        root_path=root_path,
        data_path=data_path,
        data_name=data_name)
    print('[Train_data Shape]:', train_data.shape)
    print('[Train_labels Shape]:', train_labels.shape)
    print('[Test_data Shape]:', test_data.shape)
    print('[Test_labels Shape]:', test_labels.shape)

    hardata = Dataset_HAR(
        root_path=root_path,
        data_path=data_path,
        data_name=data_name, flag='TRAIN')
    print('-[UCR data class]:', hardata)
    print('-[Num of data]', len(hardata))
    x, y = hardata[66]
    print('-[single data shape]:', x.shape)
    print('-[dim of data]:', x.shape[1])
    print('-[length of data]:', x.shape[0])
    print('-[single data label]:', y)